package SimpleLACore

import chisel3._
import chisel3.util._

class CSRFile(val CSRs : CSR*){
  def get[T <: CSR](target: T): T = CSRs.find(_.addr == target.addr).get.asInstanceOf[T]
  def read(addr: UInt) = Mux1H(CSRs.map(csr => (csr.addr.U === addr, csr.value()(this))))
  def write(addr: UInt, data: UInt, mask: UInt): Unit = CSRs.foreach(csr => {
    when(csr.addr.U === addr) {
      csr.write(data, mask)(this)
    }
  })
}

abstract class CSR extends Bundle with SimpleLACoreParam{
  def addr: Int
  def value()(implicit csrFile: CSRFile): UInt = this.asUInt
  def write(data: UInt, mask: UInt)(implicit csrFile: CSRFile): Unit = {
    this := ((data & mask) | (value() & (~mask).asUInt)).asTypeOf(this)
  }
}
object MaskWrite{
  def apply(target: UInt, data: UInt, mask: UInt): Unit = target := data & mask | target & (~mask).asUInt
}
object MaskWriteResult{
  def apply(target: UInt, data: UInt, mask: UInt): UInt = data & mask | target & (~mask).asUInt
}

class CRMD extends CSR {
  val DATM = UInt(2.W)
  val DATF = UInt(2.W)
  val PG = Bool()
  val DA = Bool()
  val IE = Bool()
  val PLV = UInt(2.W)
  
  def addr = 0x0
//  override def value()(implicit csrFile: CSRFile): UInt = this.asUInt()
//  override def write(data: UInt, mask: UInt)(implicit csrFile: CSRFile): Unit = {
//    this := ((data & mask) | (value & (~mask).asUInt)).asTypeOf(this)
//  }
}

class PRMD extends CSR {
  val PIE = Bool()
  val PPLV = UInt(2.W)
  
  def addr = 0x1
//  override def value()(implicit csrFile: CSRFile): UInt = this.asUInt()
//  override def write(data: UInt, mask: UInt)(implicit csrFile: CSRFile): Unit = {
//    this := ((data & mask) | (value & (~mask).asUInt)).asTypeOf(this)
//  }
}

class EUEN extends CSR {
  val FPE = Bool()

  def addr = 0x2
//  override def value()(implicit csrFile: CSRFile): UInt = this.asUInt()
//  override def write(data: UInt, mask: UInt)(implicit csrFile: CSRFile): Unit = {
//    this := ((data & mask) | (value & (~mask).asUInt)).asTypeOf(this)
//  }
}

class ECTL extends CSR {
  val LIE = UInt(13.W)

  def addr = 0x4
//  override def value()(implicit csrFile: CSRFile): UInt = this.asUInt()
  override def write(data: UInt, mask: UInt)(implicit csrFile: CSRFile): Unit = {
    super.write(data, mask & 0x1bff.U(13.W))
  }
}

class ESTAT extends CSR {
  val EsubCode = UInt(9.W)
  val Ecode = UInt(6.W)
  //val zero = UInt(3.W)
  val IS = Vec(13, Bool())

  def addr = 0x5
  override def value()(implicit csrFile: CSRFile): UInt = Cat(Seq(EsubCode, Ecode, 0.U(3.W)) ++ IS.reverse)
  override def write(data: UInt, mask: UInt)(implicit csrFile: CSRFile): Unit = {
    when(mask(0)){IS(0) := data(0)}
    when(mask(1)){IS(1) := data(1)}
  }
}

class ERA extends CSR {
  val PC = UInt(32.W)

  def addr = 0x6
//  override def value()(implicit csrFile: CSRFile): UInt = this.asUInt()
//  override def write(data: UInt, mask: UInt)(implicit csrFile: CSRFile): Unit = {
//    this := ((data & mask) | (value & (~mask).asUInt)).asTypeOf(this)
//  }
}

class BADV extends CSR {
  val VAddr = UInt(32.W)

  def addr = 0x7
//  override def value()(implicit csrFile: CSRFile): UInt = this.asUInt()
//  override def write(data: UInt, mask: UInt)(implicit csrFile: CSRFile): Unit = {
//    this := ((data & mask) | (value & (~mask).asUInt)).asTypeOf(this)
//  }
}

class EENTRY extends CSR {
  val VA = UInt(26.W)
  //val zero = UInt(6.W)

  def addr = 0xc
  override def value()(implicit csrFile: CSRFile): UInt = Cat(VA, 0.U(6.W))
  override def write(data: UInt, mask: UInt)(implicit csrFile: CSRFile): Unit = {
    MaskWrite(VA, data(31, 6), mask(31, 6))
//    VA := data(31, 6) & mask(31, 6) | VA & (~mask(31, 6)).asUInt
  }
}

class CPUID extends CSR {
  val CoreID = UInt(9.W)

  def addr = 0x20
//  override def value()(implicit csrFile: CSRFile): UInt = this.asUInt()
  override def write(data: UInt, mask: UInt)(implicit csrFile: CSRFile): Unit = {}
}

class SAVE0 extends CSR {
  val Data = UInt(32.W)

  def addr = 0x30
//  override def value()(implicit csrFile: CSRFile): UInt = this.asUInt()
//  override def write(data: UInt, mask: UInt)(implicit csrFile: CSRFile): Unit = {
//    this := ((data & mask) | (value & (~mask).asUInt)).asTypeOf(this)
//  }
}

class SAVE1 extends CSR {
  val Data = UInt(32.W)

  def addr = 0x31
//  override def value()(implicit csrFile: CSRFile): UInt = this.asUInt()
//  override def write(data: UInt, mask: UInt)(implicit csrFile: CSRFile): Unit = {
//    this := ((data & mask) | (value & (~mask).asUInt)).asTypeOf(this)
//  }
}

class SAVE2 extends CSR {
  val Data = UInt(32.W)

  def addr = 0x32
//  override def value()(implicit csrFile: CSRFile): UInt = this.asUInt()
//  override def write(data: UInt, mask: UInt)(implicit csrFile: CSRFile): Unit = {
//    this := ((data & mask) | (value & (~mask).asUInt)).asTypeOf(this)
//  }
}

class SAVE3 extends CSR {
  val Data = UInt(32.W)

  def addr = 0x33
//  override def value()(implicit csrFile: CSRFile): UInt = this.asUInt()
//  override def write(data: UInt, mask: UInt)(implicit csrFile: CSRFile): Unit = {
//    this := ((data & mask) | (value & (~mask).asUInt)).asTypeOf(this)
//  }
}

class LLBCTL extends CSR {
  val KLO = Bool()
//  val WCLLB = UInt(1.W)
  val ROLLB = Bool()

  def addr = 0x60
  override def value()(implicit csrFile: CSRFile): UInt = Cat(KLO, 0.U(1.W), ROLLB)
  override def write(data: UInt, mask: UInt)(implicit csrFile: CSRFile): Unit = {
    when(mask(2)){KLO := data(2)}
    when(mask(1) & data(1)){ROLLB := false.B}
  }
}

class TLBIDX extends CSR {
  val NE = Bool()
  //val zero1 = UInt(1.W)
  val PS = UInt(6.W)
  //val zero0 = UInt((24 - lgnTLB).W)
  val Index = UInt(lgnTLB.W)

  def addr = 0x10
  override def value()(implicit csrFile: CSRFile): UInt = Cat(NE, 0.U(1.W), PS, 0.U((24 - lgnTLB).W), Index)
  override def write(data: UInt, mask: UInt)(implicit csrFile: CSRFile): Unit = {
    MaskWrite(NE, data(31), mask(31))
    MaskWrite(PS, data(29, 24), mask(29, 24))
    MaskWrite(Index, data(lgnTLB - 1, 0), mask(lgnTLB - 1, 0))
  }
}

class TLBEHI extends CSR {
  val VPPN = UInt(19.W)
  //val zero = UInt(13.W)

  def addr = 0x11
  override def value()(implicit csrFile: CSRFile): UInt = Cat(VPPN, 0.U(13.W))
  override def write(data: UInt, mask: UInt)(implicit csrFile: CSRFile): Unit = {
    MaskWrite(VPPN, data(31, 13), mask(31, 13))
  }
}

class TLBELO0 extends CSR {
  val PPN = UInt(24.W)
  //val zero = UInt(1.W)
  val G = Bool()
  val MAT = UInt(2.W)
  val PLV = UInt(2.W)
  val D = Bool()
  val V = Bool()

  def addr = 0x12
  override def value()(implicit csrFile: CSRFile): UInt = Cat(PPN, 0.U(1.W), G, MAT, PLV, D, V)
  override def write(data: UInt, mask: UInt)(implicit csrFile: CSRFile): Unit = {
    super.write(data, mask)
    MaskWrite(PPN, data(31, 8), mask(31,8))
  }
}

class TLBELO1 extends CSR {
  val PPN = UInt(24.W)
  //val zero = UInt(1.W)
  val G = Bool()
  val MAT = UInt(2.W)
  val PLV = UInt(2.W)
  val D = Bool()
  val V = Bool()

  def addr = 0x13
  override def value()(implicit csrFile: CSRFile): UInt = Cat(PPN, 0.U(1.W), G, MAT, PLV, D, V)
  override def write(data: UInt, mask: UInt)(implicit csrFile: CSRFile): Unit = {
    super.write(data, mask)
    MaskWrite(PPN, data(31, 8), mask(31,8))
  }
}

class ASID extends CSR {
  val ASIDBITS = UInt(8.W)
  //val zero = UInt(6.W)
  val ASID = UInt(10.W)

  def addr = 0x18
  override def value()(implicit csrFile: CSRFile): UInt = Cat(ASIDBITS, 0.U(6.W), ASID)
  override def write(data: UInt, mask: UInt)(implicit csrFile: CSRFile): Unit = {
    MaskWrite(ASID, data(9,0), mask(9,0))
  }
}

class PGDL extends CSR {
  val Base = UInt(20.W)
  //val zero = UInt(12.W)

  def addr = 0x19
  override def value()(implicit csrFile: CSRFile): UInt = Cat(Base, 0.U(12.W))
  override def write(data: UInt, mask: UInt)(implicit csrFile: CSRFile): Unit = {
    MaskWrite(Base, data(31, 12), mask(31, 12))
  }
}

class PGDH extends CSR {
  val Base = UInt(20.W)
  //val zero = UInt(12.W)

  def addr = 0x1A
  override def value()(implicit csrFile: CSRFile): UInt = Cat(Base, 0.U(12.W))
  override def write(data: UInt, mask: UInt)(implicit csrFile: CSRFile): Unit = {
    MaskWrite(Base, data(31, 12), mask(31, 12))
  }
}

class PGD extends CSR {
//  val Base = UInt(20.W)
  //val zero = UInt(12.W)

  def addr = 0x1B
  override def value()(implicit csrFile: CSRFile): UInt =
    Mux(csrFile.get(new BADV).VAddr(31), csrFile.get(new PGDH).value(), csrFile.get(new PGDL).value())
  override def write(data: UInt, mask: UInt)(implicit csrFile: CSRFile): Unit = {}
}

class TLBRENTRY extends CSR {
  val PA = UInt(26.W)
  //val zero = UInt(6.W)

  def addr = 0x88
  override def value()(implicit csrFile: CSRFile): UInt = Cat(PA, 0.U(6.W))
  override def write(data: UInt, mask: UInt)(implicit csrFile: CSRFile): Unit = {
    MaskWrite(PA, data(31, 6), mask(31, 6))
  }
}

class DMW0 extends CSR {
  val VSEG = UInt(3.W)
  //val zero2 = UInt(1.W)
  val PSEG = UInt(3.W)
  //val zero1 = UInt(19.W)
  val MAT = UInt(2.W)
  val PLV3 = Bool()
  //val zero0 = UInt(2.W)
  val PLV0 = Bool()

  def addr = 0x180
  override def value()(implicit csrFile: CSRFile):UInt = Cat(VSEG, 0.U(1.W), PSEG, 0.U(19.W), MAT, PLV3, 0.U(2.W), PLV0)
  override def write(data: UInt, mask: UInt)(implicit csrFile: CSRFile): Unit = {
    MaskWrite(PLV0, data(0), mask(0))
    MaskWrite(PLV3, data(3), mask(3))
    MaskWrite(MAT, data(5, 4), mask(5, 4))
    MaskWrite(PSEG, data(27, 25), mask(27, 25))
    MaskWrite(VSEG, data(31, 29), mask(31, 29))
  }
}

class DMW1 extends CSR {
  val VSEG = UInt(3.W)
  //val zero2 = UInt(1.W)
  val PSEG = UInt(3.W)
  //val zero1 = UInt(19.W)
  val MAT = UInt(2.W)
  val PLV3 = Bool()
  //val zero0 = UInt(2.W)
  val PLV0 = Bool()

  def addr = 0x181
  override def value()(implicit csrFile: CSRFile):UInt = Cat(VSEG, 0.U(1.W), PSEG, 0.U(19.W), MAT, PLV3, 0.U(2.W), PLV0)
  override def write(data: UInt, mask: UInt)(implicit csrFile: CSRFile): Unit = {
    MaskWrite(PLV0, data(0), mask(0))
    MaskWrite(PLV3, data(3), mask(3))
    MaskWrite(MAT, data(5, 4), mask(5, 4))
    MaskWrite(PSEG, data(27, 25), mask(27, 25))
    MaskWrite(VSEG, data(31, 29), mask(31, 29))
  }
}

class TID extends CSR {
  val TID = UInt(32.W)

  def addr = 0x40
//  override def value()(implicit csrFile: CSRFile): UInt = this.asUInt()
//  override def write(data: UInt, mask: UInt)(implicit csrFile: CSRFile): Unit = TID := data
}

class TCFG extends CSR {
  val InitVal = UInt((TimerBits - 2).W)
  val Periodic = Bool()
  val En = Bool()

  def addr = 0x41
//  override def value()(implicit csrFile: CSRFile): UInt = this.asUInt()
  override def write(data: UInt, mask: UInt)(implicit csrFile: CSRFile): Unit = {
    val next = MaskWriteResult(value(), data, mask).asTypeOf(this)
    this := next
    when(next.En){
      csrFile.get(new TVAL).TimeVal := Cat(next.InitVal, 0.U(2.W))
      when(next.InitVal === 0.U){
        csrFile.get(new ESTAT).IS(11) := true.B
      }
    }
  }
}

class TVAL extends CSR {
  val TimeVal = UInt(TimerBits.W)

  def addr = 0x42
//  override def value()(implicit csrFile: CSRFile): UInt = this.asUInt()
  override def write(data: UInt, mask: UInt)(implicit csrFile: CSRFile): Unit = {}
}

class TICLR extends CSR {
//  val CLR = UInt(1.W)

  def addr = 0x44
//  override def value()(implicit csrFile: CSRFile): UInt = this.asUInt()
  override def write(data: UInt, mask: UInt)(implicit csrFile: CSRFile): Unit = {
    when(data(0) & mask(0)){
      csrFile.get(new ESTAT).IS(11) := false.B
    }
  }
}
